<?php
/**
 * 支持多数据库和读写分离
 *
 * Author: bluer
 * Date: 3/22/16
 */
class HyperDB extends Apdb
{
    /**
     * Mysql连接组 (dbhname => dbh)
     * @var array
     */
    protected static $_dbhs = array();

    /**
     * 可供使用的服务器组
     * @var array
     */
    protected $_servers = array('read'=>array(), 'write'=>array());

    /**
     * 记录要从主库读取数据的表(Send Reads To Masters)
     * @var array
     */
    protected $_srtm = array();

    /**
     * 当前的dbh名称
     * @var string
     */
    protected $_dbhname;

    /**
     * 当前的查询语句
     * @var string
     */
    protected $_query;

    /**
     * Master(写操作)标记
     * @var bool
     */
    protected $_useMaster = false;

    /**
     * 事务开启标记
     * @var bool
     */
    protected $_isTransaction = false;

    /**
     * 数据源调用计数
     * @var array
     */
    protected static $_dbhUseCounter = array();

    /**
     * Add the connection parameters for a database
     *
     * @param $db
     */
    function addDatabase($db) {
        extract($db, EXTR_SKIP);

        if (!isset($host) || !isset($dbname) || !isset($user) || !isset($password)) {
            return false;
        }

        isset($read)    or $read = 0;
        isset($write)   or $write = 0;

        if (!empty($host)) {
            $ident = self::getDBIdentifier($host, $dbname);
            if ($read && !array_key_exists($ident, $this->_servers)) {
                $this->_servers['read'][$ident] = $db;
            }

            if ($write && !array_key_exists($ident, $this->_servers['write'])) {
                $this->_servers['write'][$ident] = $db;
            }
        }
    }

    protected static function getDBIdentifier($host, $dbname) {
        return md5($host . $dbname);
    }

    protected static function shuffleAssoc($list) {
        if (!is_array($list)) return $list;

        $keys = array_keys($list);
        shuffle($keys);
        $random = array();
        foreach ($keys as $key) $random[$key] = $list[$key];

        return $random;
    }

    protected function checkConnection() {
        if (empty($this->_servers['write']) && empty($this->_servers['read'])) {
            $this->printError("There was no database defined.");
            return false;
        }

        // 确认当前SQL类型(读 or 写)
        $useMaster = false;
        $operator = 'read';
        $this->_dbhname = get_class($this) . '__r';
        if ($this->_useMaster || $this->isWriteQuery($this->_query)) {
            $useMaster = true;
            $this->_dbhname = get_class($this) . '__w';
            $table = $this->getTableFromQuery($this->_query);
            $this->_srtm[$table] = true;
            $operator = 'write';

        } elseif (!empty($this->_srtm)) {
            $queryMatch = substr($this->_query, 0, 5000);
            foreach ($this->_srtm as $key => $value) {
                if (false !== stripos( $queryMatch, $key )) {
                    $useMaster = true;
                    $this->_dbhname = get_class($this) . '__w';
                    $operator = 'write';
                    break;
                }
            }
        }

        // 检查是否完数据库实例
        if (array_key_exists($this->_dbhname, self::$_dbhs)
            && array_key_exists($operator, $this->_servers)
            && count($this->_servers[$operator]) > count(self::$_dbhs[$this->_dbhname])) {
            return $this->connect($useMaster);
        }

        // 检查是否有可用连接
        if (!empty(self::$_dbhs) && array_key_exists($this->_dbhname, self::$_dbhs)) {
            // 优先计数小的连接
            asort(self::$_dbhUseCounter);
            foreach (self::$_dbhUseCounter as $ident => $counter) {
                if (array_key_exists($this->_dbhname, self::$_dbhs)
                    && !array_key_exists($ident, self::$_dbhs[$this->_dbhname])) {
                    continue;
                }

                if (is_a(self::$_dbhs[$this->_dbhname][$ident], 'PDO')) {
                    $this->_dbh = self::$_dbhs[$this->_dbhname][$ident];
                    self::$_dbhUseCounter[$ident] += 1;
                    return true;
                } else {
                    unset(self::$_dbhs[$this->_dbhname][$ident]);
                    unset(self::$_dbhUseCounter[$ident]);
                }
            }
        }

        return $this->connect($useMaster);
    }

    protected function connect($useMaster) {
        // 伪随机实例列表
        if (empty(self::$_dbhs) && count($this->_servers['read']) > 1) {
            $this->_servers['read'] = $this->shuffleAssoc($this->_servers['read']);
        }

        if ($useMaster) {
            if ($this->_servers['write']) {
                // 限制仅支持单一主库
                foreach ($this->_servers['write'] as $ident => $db) {
                    if ($success = $this->_connectDB($db)) break;
                    unset($this->_servers['write'][$ident]);
                }
            }

            // 没有可写DB，清除写操作并尝试切换为只读模式
            if (!$success) {
                $this->_useMaster = false;
                $this->_srtm = array();
                $this->_query = null;
                $success = $this->checkConnection();
            }

        } else {
            if ($this->_servers['read']) {
                foreach ($this->_servers['read'] as $ident => $db) {
                    if (array_key_exists($this->_dbhname, self::$_dbhs)
                        && array_key_exists($ident, self::$_dbhs[$this->_dbhname])) {
                        continue;
                    }

                    if ($success = $this->_connectDB($db)) break;

                    unset($this->_servers['read'][$ident]);
                }
            }

            // 没有只读DB，尝试切换到主数据库
            if (!$success && empty(self::$_dbhs[$this->_dbhname])) {
                $this->_useMaster = true;
                $success = $this->checkConnection();
            } else if (!$success) {
                $success = $this->checkConnection();
                if (!$success) self::$_dbhs[$this->_dbhname] = array();
            }
        }

        return $success;
    }

    /**
     * 连接指定数据库
     * @param $db
     * @return bool
     */
    protected function _connectDB($db) {
        extract($db);

        $this->setCharset();
        try {
            $dsn = sprintf('mysql:host=%s;dbname=%s;port=%s;charset=%s',
                $host,
                $dbname,
                3306,
                $this->charset
            );

            if (empty($this->_options)) {
                $cmd = 'SET NAMES ' . $this->charset . ' COLLATE ' . $this->collate;
                $this->_options = array(PDO::ATTR_ERRMODE => PDO::ERRMODE_EXCEPTION,
                    PDO::MYSQL_ATTR_INIT_COMMAND => $cmd,
                    //PDO::ATTR_TIMEOUT =>  3
                );
            }

            if ($this->_persistent) $this->_options[PDO::ATTR_PERSISTENT] = true;

            $ident = self::getDBIdentifier($host, $dbname);
            $this->_dbh = new PDO($dsn, $user, $password, $this->_options);
            self::$_dbhs[$this->_dbhname][$ident] = $this->_dbh;
            self::$_dbhUseCounter[$ident] = 1;

        } catch(PDOException $e) {
            $this->printError($e);
            return false;
        }

        return true;
    }

    /**
     * 检查查询语句是否涉及数据改动
     * @param $q
     * @return bool
     */
    protected function isWriteQuery($q) {
        if (is_null($q) || empty($q)) return false;

        $q = ltrim($q, "\r\n\t (");
        return !preg_match('/^(?:SELECT|SHOW|DESCRIBE|DESC|EXPLAIN)\s/i', $q);
    }

    /**
     * 获取查询语句中的一个被引用的表
     *
     * @param string $q
     * @return string table
     */
    protected function getTableFromQuery($q) {
        // Remove characters that can legally trail the table name
        $q = rtrim($q, ';/-#');
        // allow (select...) union [...] style queries. Use the first queries table name.
        $q = ltrim($q, "\t (");

        // Quickly match most common queries
        if ( preg_match('/^\s*(?:'
            . 'SELECT.*?\s+FROM'
            . '|INSERT(?:\s+IGNORE)?(?:\s+INTO)?'
            . '|REPLACE(?:\s+INTO)?'
            . '|UPDATE(?:\s+IGNORE)?'
            . '|DELETE(?:\s+IGNORE)?(?:\s+FROM)?'
            . ')\s+`?([\w-]+)`?/is', $q, $maybe) )
            return $maybe[1];

        // SHOW TABLE STATUS and SHOW TABLES
        if ( preg_match('/^\s*(?:'
            . 'SHOW\s+TABLE\s+STATUS.+(?:LIKE\s+|WHERE\s+Name\s*=\s*)'
            . '|SHOW\s+(?:FULL\s+)?TABLES.+(?:LIKE\s+|WHERE\s+Name\s*=\s*)'
            . ')\W([\w-]+)\W/is', $q, $maybe) )
            return $maybe[1];

        // Big pattern for the rest of the table-related queries in MySQL 5.0
        if ( preg_match('/^\s*(?:'
            . '(?:EXPLAIN\s+(?:EXTENDED\s+)?)?SELECT.*?\s+FROM'
            . '|INSERT(?:\s+LOW_PRIORITY|\s+DELAYED|\s+HIGH_PRIORITY)?(?:\s+IGNORE)?(?:\s+INTO)?'
            . '|REPLACE(?:\s+LOW_PRIORITY|\s+DELAYED)?(?:\s+INTO)?'
            . '|UPDATE(?:\s+LOW_PRIORITY)?(?:\s+IGNORE)?'
            . '|DELETE(?:\s+LOW_PRIORITY|\s+QUICK|\s+IGNORE)*(?:\s+FROM)?'
            . '|DESCRIBE|DESC|EXPLAIN|HANDLER'
            . '|(?:LOCK|UNLOCK)\s+TABLE(?:S)?'
            . '|(?:RENAME|OPTIMIZE|BACKUP|RESTORE|CHECK|CHECKSUM|ANALYZE|OPTIMIZE|REPAIR).*\s+TABLE'
            . '|TRUNCATE(?:\s+TABLE)?'
            . '|CREATE(?:\s+TEMPORARY)?\s+TABLE(?:\s+IF\s+NOT\s+EXISTS)?'
            . '|ALTER(?:\s+IGNORE)?\s+TABLE'
            . '|DROP\s+TABLE(?:\s+IF\s+EXISTS)?'
            . '|CREATE(?:\s+\w+)?\s+INDEX.*\s+ON'
            . '|DROP\s+INDEX.*\s+ON'
            . '|LOAD\s+DATA.*INFILE.*INTO\s+TABLE'
            . '|(?:GRANT|REVOKE).*ON\s+TABLE'
            . '|SHOW\s+(?:.*FROM|.*TABLE)'
            . ')\s+`?([\w-]+)`?/is', $q, $maybe) )
            return $maybe[1];
    }

    protected function _do_query($query, $bind = null) {
        $this->_query = $query;

        if (!$this->checkConnection()) {
            $this->show_errors = true;
            $this->printError("Establishing a database connection failed.");
            return;
        }

        // 无可写数据库可用时，取消写入操作
        if (is_null($this->_query)) {
            $this->_result = false;
            return;
        }

        if (defined('SAVEQUERIES') && SAVEQUERIES) {
            $time_start = $this->timer_start();
        }

        try {
            if (!is_null($bind) && is_array($bind)) {
                $this->_stmt = $this->_dbh->prepare($query);
                $this->bind($bind);

                if ($this->_stmt->execute()) {
                    $this->_result = true;
                } else {
                    $this->_result = false;
                }
            } else {
                $this->_stmt = $this->_dbh->query($query);
                $this->_result = true;
            }
            $this->_num_queries++;

            if (defined('SAVEQUERIES') && SAVEQUERIES) {
                $this->_queries[] = array($query, $this->timer_stop($time_start));
            }

        } catch (PDOException $e) {
            $this->_result = false;
            $this->printError($e);
        }
    }

    public function prepare($query, $args) {
        if (is_null($query)) return null;

        if (!$this->checkConnection()) {
            $this->show_errors = true;
            $this->printError("Establishing a database connection failed.");
            return;
        }

        $args = func_get_args();
        array_shift($args);
        // If args were passed as an array (as in vsprintf), move them up
        if (isset($args[0]) && is_array($args[0]))
            $args = $args[0];

        // in case someone mistakenly already singlequoted it
        $query = str_replace("'%s'", '%s', $query);

        // doublequote unquoting
        $query = str_replace('"%s"', '%s', $query);

        // Force floats to be locale unaware
        $query = preg_replace('|(?<!%)%f|' , '%F', $query);
        foreach ($args as &$arg) {
            if (!is_int($arg))
                $arg = $this->_dbh->quote($arg);
        }

        return @vsprintf($query, $args);
    }

    public function beginTransaction() {
        $this->_useMaster = true;
        $this->_isTransaction = true;

        if ($this->checkConnection() && $this->_useMaster) {
            return $this->_dbh->beginTransaction();
        }

        return false;
    }

    public function endTransaction() {
        $result = false;

        if ($this->checkConnection() && $this->_useMaster) {
            $result = $this->_dbh->commit();
            $this->_useMaster = false;
            $this->_isTransaction = false;
        }

        return $result;
    }

    public function cancelTransaction() {
        $result = false;
        if ($this->checkConnection() && $this->_useMaster) {
            $result = $this->_dbh->rollBack();
            $this->_useMaster = false;
            $this->_isTransaction = false;
        }

        return $result;
    }

    public function lastInsertId() {
        $id = 0;
        $this->_useMaster = true;
        if ($this->checkConnection() && $this->_useMaster) {
            $id = $this->_dbh->lastInsertId();
            if (!$this->_isTransaction) $this->_useMaster = false;
        }

        return $id;
    }
}

?>